#include "unitree_arm_sdk/control/unitreeArm.h"
#include <iostream>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <cstring>
#include <iomanip>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <atomic>
#include <fcntl.h>
#include <sys/select.h>
#include <fstream>
#include <chrono>
#include <ctime>
#include <sstream>
#include <string>
#include <cerrno>
#include <termios.h>
#include "unitree_arm_sdk/thirdparty/robotics.h" // 假设这里有rpyxyz到HomoMat的工具函数
#include "unitree_arm_sdk/model/ArmModel.h"      // 包含ArmModel、Z1Model声明
#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <termios.h>
#include <fcntl.h>
#include <unistd.h>
#include <iomanip>
#include <cstring>
// #define USE_SPEED_CONTROL
#define USE_POSITION_CONTROL
#define thresh 0.2
// 在这里定义USE_PREDICTIVE_CONTROL来控制是否启用预测控制
// 注释掉下面这一行则不使用预测控制
// #define USE_PREDICTIVE_CONTROL

using namespace UNITREE_ARM;

// ===== 改为“只处理最新数据”，故不用顺序队列了 =====
// 全局仅存一个“最新packet”和互斥量
static std::mutex g_latestMutex;
static std::vector<unsigned char> g_latestPacket;
static std::atomic<bool> g_hasNewPacket(false);

float force[6] = {0.0f}; // 全局力数据
std::mutex force_mutex;
std::condition_variable force_cv;
std::atomic<bool> exit_flag(false);
std::mutex arm_mutex;

uint16_t crc16_modbus(const uint8_t *data, size_t len)
{
    uint16_t crc = 0xFFFF;
    for (size_t i = 0; i < len; i++)
    {
        crc ^= (uint16_t)data[i];
        for (int j = 0; j < 8; j++)
        {
            if (crc & 0x0001)
            {
                crc >>= 1;
                crc ^= 0xA001;
            }
            else
            {
                crc >>= 1;
            }
        }
    }
    return crc;
}
using namespace std;
void print_header_check(const unsigned char *response, size_t len)
{
    // Print the first two bytes in Hex format
    if (len >= 2)
    {
        std::cout << "First 2 bytes: ";
        for (int i = 0; i < 2; i++)
        {
            std::cout << hex << uppercase << setw(2) << setfill('0') << (int)response[i] << " ";
        }
        cout << "(Hex)" << endl;
        cout << "Response Length: " << len << endl;
    }
}

void parse_response(const unsigned char *response, size_t len)
{
    // Check if response length is valid
    if (len < 27)
    {
        cout << "Invalid response length." << endl;
        return;
    }

    print_header_check(response, len); // Print first two bytes for debugging

    // Header check (0xAA, 0x55)
    if (response[0] != 0xAA || response[1] != 0x55)
    {
        cout << "Invalid header." << endl;
        return;
    }

    // Get package length (2 bytes)
    unsigned int package_length = (response[2] << 8) | response[3];
    // cout << "Package Length: " << package_length << " bytes" << endl;

    // Get package number (2 bytes)
    unsigned int package_no = (response[4] << 8) | response[5];
    // cout << "Package No: " << package_no << endl;

    // Get the data part (24 bytes)
    unsigned char data[24];
    memcpy(data, response + 6, 24);
    // cout << "Data (Hex): ";
    // for (int i = 0; i < 24; i++) {
    //     cout << hex << uppercase << setw(2) << setfill('0') << (int)data[i] << " ";
    // }
    // cout << endl;

    // Convert each channel data (4 bytes per channel) to a float
    for (int i = 0; i < 6; i++)
    {
        float value;
        memcpy(&value, data + i * 4, 4);
        // cout << "Channel " << i + 1 << " (Engineering Unit): " << fixed << setprecision(6) << value << " N (Force) / Nm (Torque)" << endl;
        force[i] = value;
    }

    // Get the checksum (last byte)
    unsigned char checksum = response[len - 1];
    // cout << "SUM Checksum: " << hex << uppercase << setw(2) << setfill('0') << (int)checksum << endl;
}

int initSerialPort(const std::string &portname)
{
    int fd = open(portname.c_str(), O_RDWR | O_NOCTTY | O_SYNC);
    if (fd < 0)
    {
        std::cerr << "[initSerialPort] Error opening " << portname << ": " << strerror(errno) << std::endl;
        return -1;
    }

    struct termios tty;
    memset(&tty, 0, sizeof tty);
    if (tcgetattr(fd, &tty) != 0)
    {
        std::cerr << "[initSerialPort] Error from tcgetattr: " << strerror(errno) << std::endl;
        close(fd);
        return -1;
    }

    cfsetospeed(&tty, B115200);
    cfsetispeed(&tty, B115200);

    tty.c_cflag = (tty.c_cflag & ~CSIZE) | CS8;
    tty.c_iflag &= ~IGNBRK;
    tty.c_lflag = 0;
    tty.c_oflag = 0;
    tty.c_cc[VMIN] = 0;
    tty.c_cc[VTIME] = 5;
    tty.c_iflag &= ~(IXON | IXOFF | IXANY);
    tty.c_cflag |= (CLOCAL | CREAD);
    tty.c_cflag &= ~(PARENB | PARODD);
    tty.c_cflag &= ~CSTOPB;
    tty.c_cflag &= ~CRTSCTS;

    if (tcsetattr(fd, TCSANOW, &tty) != 0)
    {
        std::cerr << "[initSerialPort] Error from tcsetattr: " << strerror(errno) << std::endl;
        close(fd);
        return -1;
    }
    return fd;
}

class Z1ARM : public unitreeArm
{
public:
    Z1ARM() : unitreeArm(true), tcp_sockfd(-1) {};
    ~Z1ARM()
    {
        if (tcp_sockfd != -1)
        {
            close(tcp_sockfd);
            std::cout << "[~Z1ARM] TCP socket closed." << std::endl;
        }
    };

    void armCtrlByFSM();
    void armCtrlInJointCtrl();
    void armCtrlInCartesian();
    void printState();

    double gripper_pos = 0.0;
    double joint_speed = 2.0;
    double cartesian_speed = 3.0;

    void setForwardPosition(const Vec6 &pos) { forward_position = pos; }
    Vec6 getForwardPosition() const { return forward_position; }

    int tcp_sockfd;

private:
    Vec6 forward_position;
};

void Z1ARM::armCtrlByFSM()
{
    std::cout << "[armCtrlByFSM] Before labelRun(\"forward\")" << std::endl;
    try
    {
        labelRun("forward");
        std::cout << "[armCtrlByFSM] labelRun(\"forward\") OK." << std::endl;
    }
    catch (...)
    {
        std::cerr << "[armCtrlByFSM] Exception in labelRun." << std::endl;
        return;
    }

    std::cout << "[armCtrlByFSM] After labelRun" << std::endl;
    printState();

    {
        std::lock_guard<std::mutex> lock(arm_mutex);
        forward_position = lowstate->endPosture;
        std::cout << "[armCtrlByFSM] Save forward_position: " << forward_position.transpose() << std::endl;
    }
}

void Z1ARM::armCtrlInJointCtrl()
{
    std::cout << "[armCtrlInJointCtrl] Before labelRun(\"forward\")" << std::endl;
    try
    {
        labelRun("forward");
        std::cout << "[armCtrlInJointCtrl] labelRun(\"forward\") OK." << std::endl;
    }
    catch (...)
    {
        std::cerr << "[armCtrlInJointCtrl] Exception in labelRun." << std::endl;
        return;
    }

    std::cout << "[armCtrlInJointCtrl] After labelRun" << std::endl;
    startTrack(ArmFSMState::JOINTCTRL);

    Timer timer(_ctrlComp->dt);
    for (int i = 0; i < 1000; i++)
    {
        if (exit_flag)
        {
            std::cout << "[armCtrlInJointCtrl] Exit flag set, break." << std::endl;
            break;
        }
        directions << 0, 0, 0, -1, 0, 0, -1;
        joint_speed = 1.0;
        jointCtrlCmd(directions, joint_speed);
        timer.sleep();
    }
    std::cout << "[armCtrlInJointCtrl] Completed." << std::endl;
}

void Z1ARM::armCtrlInCartesian()
{
    std::cout << "[armCtrlInCartesian] Before labelRun(\"forward\")" << std::endl;
    try
    {
        labelRun("forward");
        std::cout << "[armCtrlInCartesian] labelRun(\"forward\") OK." << std::endl;
    }
    catch (...)
    {
        std::cerr << "[armCtrlInCartesian] Exception in labelRun." << std::endl;
        return;
    }

    std::cout << "[armCtrlInCartesian] After labelRun" << std::endl;
    startTrack(ArmFSMState::CARTESIAN);

    double angular_vel = 0.3;
    double linear_vel = 0.3;
    Timer timer(_ctrlComp->dt);
    for (int i = 0; i < 2000; i++)
    {
        if (exit_flag)
        {
            std::cout << "[armCtrlInCartesian] Exit flag set, break." << std::endl;
            break;
        }
        directions << 0, 0, 0, 0, 0, -1, -1;
        cartesianCtrlCmd(directions, angular_vel, linear_vel);
        timer.sleep();
    }
    std::cout << "[armCtrlInCartesian] Completed." << std::endl;
}

void Z1ARM::printState()
{
    std::cout << "------ joint State ------" << std::endl;
    std::cout << "qState: " << lowstate->getQ().transpose() << std::endl;
    std::cout << "qdState: " << lowstate->getQd().transpose() << std::endl;
    std::cout << "tauState: " << lowstate->getTau().transpose() << std::endl;
    std::cout << "------ Endeffector Cartesian Posture ------" << std::endl;
    std::cout << "(roll pitch yaw x y z) = " << lowstate->endPosture.transpose() << std::endl;
}

void monitorKeyboard(std::atomic<bool> &exit_flag)
{
    std::cout << "[monitorKeyboard] Press 'q' followed by Enter to exit." << std::endl;
    char c;
    while (std::cin >> c)
    {
        if (c == 'q' || c == 'Q')
        {
            exit_flag = true;
            std::cout << "[monitorKeyboard] 'q' pressed. Exit..." << std::endl;
            force_cv.notify_all();
            break;
        }
    }
}

void tcpClient(Z1ARM &arm)
{
    arm.tcp_sockfd = socket(AF_INET, SOCK_STREAM, 0);
    if (arm.tcp_sockfd < 0)
    {
        std::cerr << "[tcpClient] Failed to create socket: " << strerror(errno) << std::endl;
        return;
    }

    struct sockaddr_in server_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(14532);
    server_addr.sin_addr.s_addr = inet_addr("115.236.153.174");

    std::cout << "[tcpClient] Attempt to connect..." << std::endl;

    fd_set set;
    FD_ZERO(&set);
    FD_SET(arm.tcp_sockfd, &set);

    struct timeval timeout;
    timeout.tv_sec = 5;
    timeout.tv_usec = 0;

    int ret = connect(arm.tcp_sockfd, (struct sockaddr *)&server_addr, sizeof(server_addr));
    if (ret < 0)
    {
        if (errno == EINPROGRESS)
        {
            ret = select(arm.tcp_sockfd + 1, NULL, &set, NULL, &timeout);
            if (ret <= 0)
            {
                std::cerr << "[tcpClient] Connect timed out or err: " << strerror(errno) << std::endl;
                close(arm.tcp_sockfd);
                arm.tcp_sockfd = -1;
                return;
            }
            else
            {
                int so_error = 0;
                socklen_t len = sizeof(so_error);
                getsockopt(arm.tcp_sockfd, SOL_SOCKET, SO_ERROR, &so_error, &len);
                if (so_error != 0)
                {
                    std::cerr << "[tcpClient] Connect fail so_error=" << so_error << ", " << strerror(so_error) << std::endl;
                    close(arm.tcp_sockfd);
                    arm.tcp_sockfd = -1;
                    return;
                }
            }
        }
        else
        {
            std::cerr << "[tcpClient] Connect fail: " << strerror(errno) << std::endl;
            close(arm.tcp_sockfd);
            arm.tcp_sockfd = -1;
            return;
        }
    }
    std::cout << "[tcpClient] Connected ok." << std::endl;
}

void sendForceData(Z1ARM &arm, float *force)
{
    if (arm.tcp_sockfd == -1)
    {
        std::cerr << "[sendForceData] TCP not connected." << std::endl;
        return;
    }

    uint8_t packet[4 + 4 * 6 + 2]; // Adjusted to include all force channels (6 values)
    packet[0] = 0xAA;
    packet[1] = 0xBB;
    packet[2] = 0xCC;
    packet[3] = 0xDD;

    // Debug: print force values
    // std::cout << "[sendForceData] Force values: ";
    // for (int i = 0; i < 6; i++) {
    //     std::cout << force[i] << " ";
    // }
    std::cout << std::endl;

    std::memcpy(&packet[4], force, sizeof(float) * 6); // Copy the force array to the packet

    uint16_t crc = crc16_modbus(packet, 4 + 6 * 4);   // Adjusted to include the entire force data
    packet[4 + 6 * 4] = (uint8_t)(crc & 0xFF);        // CRC low byte
    packet[5 + 6 * 4] = (uint8_t)((crc >> 8) & 0xFF); // CRC high byte

    // Debug: print packet before sending
    // std::cout << "[sendForceData] Sending packet: ";
    for (int i = 0; i < sizeof(packet); i++)
    {
        // std::cout << std::hex << std::uppercase << (int)packet[i] << " ";
    }
    std::cout << std::dec << std::endl;

    ssize_t sent = send(arm.tcp_sockfd, (char *)packet, sizeof(packet), 0);
    if (sent < 0)
    {
        std::cerr << "[sendForceData] Failed: " << strerror(errno) << std::endl;
    }
    else
    {
        // std::cout << "[sendForceData] Sent " << sent << " bytes." << std::endl;
    }
}

void logArmPosition(Z1ARM &arm, std::atomic<bool> &exit_flag, const std::string &filename)
{
    std::ofstream outfile(filename, std::ios::out);
    if (!outfile.is_open())
    {
        std::cerr << "[logArmPosition] Cannot open file: " << filename << std::endl;
        return;
    }

    std::cout << "[logArmPosition] Logging posture to " << filename << std::endl;
    outfile << "timestamp_ms roll pitch yaw x y z Fx Fy Fz Tx Ty Tz q0 q1 q2 q3 q4 q5 qd0 qd1 qd2 qd3 qd4 qd5 tau0 tau1 tau2 tau3 tau4 tau5\n";
    outfile.flush();

    std::chrono::milliseconds interval(10);
    while (!exit_flag)
    {
        auto now = std::chrono::system_clock::now();
        auto now_ms = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch()) % 1000;

        auto now_time_t = std::chrono::system_clock::to_time_t(now);
        std::tm tm_time;
        localtime_r(&now_time_t, &tm_time);

        std::ostringstream timestamp_stream;
        timestamp_stream << std::put_time(&tm_time, "%Y-%m-%d_%H-%M-%S") << '.' << std::setfill('0') << std::setw(3) << now_ms.count();

        Vec6 endPosture;
        {
            std::lock_guard<std::mutex> lk(arm_mutex);
            if (arm.lowstate == nullptr)
            {
                std::cerr << "[logArmPosition] arm.lowstate=nullptr\n";
                endPosture = Vec6::Zero();
            }
            else
            {
                endPosture = arm.lowstate->endPosture;
            }
        }

        Vec6 Q = arm.lowstate->getQ();
        Vec6 Qd = arm.lowstate->getQd();
        Vec6 Qt = arm.lowstate->getTau();

        std::lock_guard<std::mutex> lk(force_mutex);
        float force_vals[6];
        std::copy(std::begin(force), std::end(force), std::begin(force_vals));

        outfile << timestamp_stream.str() << " "
                << endPosture[0] << " " << endPosture[1] << " " << endPosture[2] << " "
                << endPosture[3] << " " << endPosture[4] << " " << endPosture[5] << " "
                << force[0] << " " << force[1] << " " << force[2] << " "
                << force[3] << " " << force[4] << " " << force[5] << " "
                << Q[0] << " " << Q[1] << " " << Q[2] << " " << Q[3] << " " << Q[4] << " " << Q[5] << " "
                << Qd[0] << " " << Qd[1] << " " << Qd[2] << " " << Qd[3] << " " << Qd[4] << " " << Qd[5] << " "
                << Qt[0] << " " << Qt[1] << " " << Qt[2] << " " << Qt[3] << " " << Qt[4] << " " << Qt[5] << "\n";
        outfile.flush();

        std::this_thread::sleep_for(interval);
    }

    outfile.close();
    std::cout << "[logArmPosition] Logging stopped.\n";
}

// ============= receiveDataFromServer: 只存最新数据 (no queue), + dynamic delay calc
void receiveDataFromServer(Z1ARM &arm)
{
    // 改为“只保留最新帧” - 不再用队列
    // 并计算并存储[40..43] = delay_ms
    std::vector<unsigned char> buffer;
    const int PACKET_SIZE = 40;

    int flags = fcntl(arm.tcp_sockfd, F_GETFL, 0);
    if (flags == -1)
    {
        std::cerr << "[receiveDataFromServer] fcntl get err.\n";
        exit_flag = true;
        return;
    }
    if (fcntl(arm.tcp_sockfd, F_SETFL, flags | O_NONBLOCK) == -1)
    {
        std::cerr << "[receiveDataFromServer] fcntl set O_NONBLOCK err.\n";
        exit_flag = true;
        return;
    }

    while (!exit_flag)
    {
        unsigned char tmp[512];
        ssize_t bytes = recv(arm.tcp_sockfd, (char *)tmp, sizeof(tmp), 0);
        if (bytes < 0)
        {
            if (errno == EWOULDBLOCK || errno == EAGAIN)
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(100));
                continue;
            }
            else
            {
                std::cerr << "[receiveDataFromServer] recv err:" << strerror(errno) << "\n";
                exit_flag = true;
                break;
            }
        }
        else if (bytes == 0)
        {
            std::cerr << "[receiveDataFromServer] remote closed.\n";
            exit_flag = true;
            break;
        }
        buffer.insert(buffer.end(), tmp, tmp + bytes);

        while (buffer.size() >= PACKET_SIZE && !exit_flag)
        {
            size_t pos = buffer.size();
            bool found_header = false;
            for (size_t i = 0; i <= buffer.size() - PACKET_SIZE; i++)
            {
                if (buffer[i] == 0xAA && buffer[i + 1] == 0xBB &&
                    buffer[i + 2] == 0xCC && buffer[i + 3] == 0xDD)
                {
                    pos = i;
                    found_header = true;
                    break;
                }
            }
            if (!found_header)
            {
                buffer.clear();
                break;
            }

            if (buffer.size() - pos >= PACKET_SIZE)
            {
                std::vector<unsigned char> packet(buffer.begin() + pos, buffer.begin() + pos + PACKET_SIZE);
                if (packet[38] == 0xEE && packet[39] == 0xFF)
                {
                    // parse
                    float pos_array[7]; // not necessarily used
                    memcpy(pos_array, &packet[4], 7 * sizeof(float));
                    uint32_t send_ts = 0;
                    memcpy(&send_ts, &packet[32], 4);
                    uint16_t cnt = (packet[36] << 8) | (packet[37]);

                    // compute local_ms
                    struct timeval tv;
                    gettimeofday(&tv, nullptr);
                    time_t raw_time = tv.tv_sec;
                    struct tm local_tm;
                    localtime_r(&raw_time, &local_tm);

                    uint32_t local_ms = (uint32_t)local_tm.tm_hour * 3600000U + (uint32_t)local_tm.tm_min * 60000U + (uint32_t)local_tm.tm_sec * 1000U + (uint32_t)(tv.tv_usec / 1000U);

                    int32_t delay_ms = (int32_t)(local_ms) - (int32_t)(send_ts);
                    // store [40..43]
                    packet.resize(44);
                    memcpy(&packet[40], &delay_ms, sizeof(int32_t));

                    // 只存最新
                    {
                        std::lock_guard<std::mutex> lk(g_latestMutex);
                        g_latestPacket = packet;
                        g_hasNewPacket.store(true);
                    }

                    buffer.erase(buffer.begin(), buffer.begin() + pos + PACKET_SIZE);
                }
                else
                {
                    buffer.erase(buffer.begin(), buffer.begin() + pos + 4);
                }
            }
            else
            {
                break;
            }
        }
    }
}

void armControlThreadFunc(Z1ARM &arm)
{
#ifdef USE_POSITION_CONTROL
    {
        std::lock_guard<std::mutex> lk(arm_mutex);
        arm.startTrack(ArmFSMState::JOINTCTRL);
    }
#endif

    // 新增的静态变量或全局变量，用于实现微调模式
    // 如果你希望在函数内多次进入/退出，也可将这些变量设为static，或在更高层管理
    static bool microAdjustActive = false; // 标记是否处于微调模式
    static Vec6 microAdjustBase;           // 记录进入微调模式时的机械臂基准位置
    static Vec6 init_pos;
    while (!exit_flag)
    {
        // 若没有新数据就 sleep
        if (!g_hasNewPacket.load())
        {
            std::this_thread::sleep_for(std::chrono::milliseconds(50));
            continue;
        }

        // 取最新数据
        std::vector<unsigned char> raw_packet;
        {
            std::lock_guard<std::mutex> lk(g_latestMutex);
            raw_packet = g_latestPacket;
            g_hasNewPacket.store(false);
        }

        // 数据长度检查
        if (raw_packet.size() < 44)
        {
            continue;
        }

        float pos_array[7];
        double roll = 0, pitch = 0, yaw = 0;
        memcpy(pos_array, &raw_packet[4], 7 * sizeof(float));
        float x = pos_array[0];
        float y = pos_array[1];
        float z = pos_array[2];
        roll = pos_array[5];
        pitch = pos_array[3];
        yaw = pos_array[4];

        std::cout << "[armControlThreadFunc] x=" << x << ",y=" << y << ",z=" << z
                  << ",pitch=" << pitch << ",yaw=" << yaw << ",roll=" << roll << "\n";

        // 判断 pos_array[6] 是否大于阈值，若是则进入或保持微调模式
        if (pos_array[6] > thresh)
        {
            // 如果还没进入过微调模式，则记录当前机械臂位置作为基准
            if (!microAdjustActive)
            {
                {
                    std::lock_guard<std::mutex> lk(arm_mutex);
                    microAdjustBase = arm.lowstate->endPosture; // 记录当下姿态
                    init_pos[0] = pos_array[0];
                    init_pos[1] = pos_array[1];
                    init_pos[2] = pos_array[2];
                }
                microAdjustActive = true;
                std::cout << "[armControlThreadFunc] >>> Enter micro-adjust mode! <<<\n";
            }

            // 处于微调模式时：在 microAdjustBase 的基础上累加增量
            Vec6 new_position;
            {
                std::lock_guard<std::mutex> lk(arm_mutex);

                // 以基准位置为起点
                new_position = microAdjustBase;

                // 对应 xyz 增量
                new_position[3] = microAdjustBase[3] + (1 - pos_array[6]) * (z - init_pos[2]) / 450.0f;
                new_position[4] = microAdjustBase[4] + (1 - pos_array[6]) * (x - init_pos[0]) / 900.0f;
                new_position[5] = microAdjustBase[5] + (1 - pos_array[6]) * (y - init_pos[1]) / 900.0f;

                // 对应 roll pitch yaw 增量
                new_position[0] = roll / 1.0f;
                new_position[1] = pitch / 10.0f;
                new_position[2] = yaw / 10.0f;
            }

            // 下面与原先做 IK、插值、下发命令的逻辑基本相同
            RotMat R = robo::rpyToRotMat(new_position[0], new_position[1], new_position[2]);
            Eigen::Vector3d p(new_position[3], new_position[4], new_position[5]);
            Eigen::MatrixXd Ttemp = robo::RpToTrans(R, p);
            HomoMat TDes = Ttemp;

#ifdef USE_POSITION_CONTROL
            Vec6 qPast, qdPast, q_result;
            {
                std::lock_guard<std::mutex> lk(arm_mutex);
                qPast = arm.lowstate->getQ();
                qdPast = arm.lowstate->getQd();
            }

            auto armModel = dynamic_cast<Z1Model *>(arm._ctrlComp->armModel);
            if (!armModel)
            {
                std::cerr << "[armControlThreadFunc] armModel is not Z1Model.\n";
                continue;
            }

            // 这里视需求决定是否启用预测
#ifdef USE_PREDICTIVE_CONTROL
            int32_t delay_ms = 0;
            memcpy(&delay_ms, &raw_packet[40], sizeof(int32_t));
            double dynamic_latency = delay_ms / 1000.0;
            std::cout << "[armControlThreadFunc] dynamic_latency=" << dynamic_latency << " s.\n";

            Vec6 qFuture = qPast + qdPast * dynamic_latency;
            std::cout << " qFuture= " << qFuture.transpose() << "\n";

            if (!armModel->inverseKinematics(TDes, qFuture, q_result, true))
            {
                std::cerr << "IK failed with predictive.\n";
                continue;
            }
#else
            // 不启用预测
            if (!armModel->inverseKinematics(TDes, qPast, q_result, true))
            {
                std::cerr << "IK fail.\n";
                continue;
            }
#endif
            {
                std::lock_guard<std::mutex> lk(arm_mutex);
                Vec6 currentQ = arm.lowstate->getQ();
                int duration = 100;
                Timer timer(arm._ctrlComp->dt);
                Vec6 qd = (q_result - currentQ) / (duration * arm._ctrlComp->dt);

                for (int i = 0; i < duration && !exit_flag; i++)
                {
                    double t = (double)i / duration;
                    arm.q = currentQ * (1 - t) + q_result * t;
                    arm.qd = qd;
                    arm.setArmCmd(arm.q, arm.qd);
                    timer.sleep();
                }
                arm.printState();
            }
#endif // USE_POSITION_CONTROL
        }
        else
        {
            // pos_array[6] <= thresh，退出微调模式(若已在微调模式则退出)
            if (microAdjustActive)
            {
                microAdjustActive = false;
                std::cout << "[armControlThreadFunc] <<< Exit micro-adjust mode! >>>\n";
            }
            // 仍采用原有的控制逻辑：直接将姿态映射到 new_position
            Vec6 new_position;
            {
                std::lock_guard<std::mutex> lk(arm_mutex);
                new_position = arm.lowstate->endPosture;
                new_position[3] = (z - 0) / 450.0f;
                new_position[4] = x / 900.0f;
                new_position[5] = 0.31 + y / 900.0f;
                /******************拟合后映射公式************* */
                new_position[3] = 4.9226 * (z / 1000.0f + 0.0073) - 0.6789;
                new_position[4] = 5.3587 * (x / 1000.0f + 0.1272) - 0.6801;
                new_position[5] = 3.8276 * (y / 1000.0f + 0.1272) - 0.1767;
                /************************************ */
                new_position[0] = roll / 1.0f;
                new_position[1] = pitch / 10.0f;
                new_position[2] = yaw / 10.0f;
            }
            cout << "target: " << new_position[0] << " " << new_position[1] << " " << new_position[2] << " " << new_position[3] << " " << new_position[4] << " " << new_position[5] << std::endl;
            RotMat R = robo::rpyToRotMat(new_position[0], new_position[1], new_position[2]);
            Eigen::Vector3d p(new_position[3], new_position[4], new_position[5]);
            Eigen::MatrixXd Ttemp = robo::RpToTrans(R, p);
            HomoMat TDes = Ttemp;

#ifdef USE_POSITION_CONTROL
            Vec6 qPast, qdPast, q_result;
            {
                std::lock_guard<std::mutex> lk(arm_mutex);
                qPast = arm.lowstate->getQ();
                qdPast = arm.lowstate->getQd();
            }

            auto armModel = dynamic_cast<Z1Model *>(arm._ctrlComp->armModel);
            if (!armModel)
            {
                std::cerr << "[armControlThreadFunc] armModel is not Z1Model.\n";
                continue;
            }

#ifdef USE_PREDICTIVE_CONTROL
            int32_t delay_ms = 0;
            memcpy(&delay_ms, &raw_packet[40], sizeof(int32_t));
            double dynamic_latency = delay_ms / 1000.0;
            std::cout << "[armControlThreadFunc] dynamic_latency=" << dynamic_latency << " s.\n";

            Vec6 qFuture = qPast + qdPast * dynamic_latency;
            std::cout << " qFuture= " << qFuture.transpose() << "\n";

            if (!armModel->inverseKinematics(TDes, qFuture, q_result, true))
            {
                std::cerr << "IK failed with predictive.\n";
                continue;
            }
#else
            if (!armModel->inverseKinematics(TDes, qPast, q_result, true))
            {
                std::cerr << "IK fail.\n";
                continue;
            }
#endif
            {
                std::lock_guard<std::mutex> lk(arm_mutex);
                Vec6 currentQ = arm.lowstate->getQ();
                int duration = 100;
                Timer timer(arm._ctrlComp->dt);
                Vec6 qd = (q_result - currentQ) / (duration * arm._ctrlComp->dt);

                for (int i = 0; i < duration && !exit_flag; i++)
                {
                    double t = (double)i / duration;
                    arm.q = currentQ * (1 - t) + q_result * t;
                    arm.qd = qd;
                    arm.setArmCmd(arm.q, arm.qd);
                    timer.sleep();
                }
                arm.printState();
            }
#endif // USE_POSITION_CONTROL
        }
    }
}

void armSpeedControlThreadFunc(Z1ARM &arm)
{
    {
        std::lock_guard<std::mutex> lk(arm_mutex);
        arm.startTrack(ArmFSMState::CARTESIAN);
    }

    double angular_vel = 0.3, linear_vel = 0.3;

    while (!exit_flag)
    {
        if (!g_hasNewPacket.load())
        {
            std::this_thread::sleep_for(std::chrono::milliseconds(50));
            continue;
        }
        std::vector<unsigned char> raw_packet;
        {
            std::lock_guard<std::mutex> lk(g_latestMutex);
            raw_packet = g_latestPacket;
            g_hasNewPacket.store(false);
        }
        if (raw_packet.size() < 40)
            continue;

        float px, py, pz;
        memcpy(&px, &raw_packet[4], sizeof(float));
        memcpy(&py, &raw_packet[8], sizeof(float));
        memcpy(&pz, &raw_packet[12], sizeof(float));

        Vec7 directions;
        directions << 0, 0, 0, (pz - 100) > 0 ? 0.5 : -0.5, (py > 0) ? 0.5 : -0.5, (px > 0) ? 0.5 : -0.5, 0;

        {
            std::lock_guard<std::mutex> lk(arm_mutex);
            std::cout << "[armSpeedControlThreadFunc] directions=" << directions.transpose() << "\n";
            arm.cartesianCtrlCmd(directions, angular_vel, linear_vel);
            arm.printState();
        }
    }
}
unsigned char cmd[] = "AT+GOD\r\n";
unsigned char response[100]; // Buffer for the response
void readSerialData(int fd, std::atomic<bool> &exit_flag)
{
    std::string buffer;
    char c;
    while (!exit_flag)
    {
        tcflush(fd, TCIFLUSH);
        write(fd, cmd, sizeof(cmd) - 1);
        usleep(1000000);
        int bytes_read = read(fd, response, sizeof(response));
        if (bytes_read > 0)
        {
            print_header_check(response, bytes_read); // Print first 2 bytes for debugging

            // Check header
            if (response[0] == 0xAA && response[1] == 0x55)
            {
                parse_response(response, bytes_read); // Process valid frame
            }
            else
            {
                cout << "Invalid header." << endl;
            }
        }
        else
        {
            cout << "No data received." << endl;
        }
        usleep(1000000); // Wait before sending the next command
    }
    close(fd);
}

void sendForceThreadFunc(Z1ARM &arm)
{
    while (!exit_flag)
    {
        float force_vals[6];
        {
            std::unique_lock<std::mutex> lk(force_mutex);
            std::copy(std::begin(force), std::end(force), std::begin(force_vals));
        }
        // std::cout << "[sendForceThreadFunc] Force data to send: ";
        for (int i = 0; i < 6; i++)
        {
            std::cout << force_vals[i] << " ";
        }
        std::cout << std::endl;
        sendForceData(arm, force_vals);
        usleep(100000); // Sleep for 100 ms
    }
}

int main()
{
    std::cout << std::fixed << std::setprecision(3);

    int serial_fd = initSerialPort("/dev/ttyUSB0");
    Z1ARM arm;
    if (arm.sendRecvThread != nullptr)
    {
        arm.sendRecvThread->start();
        std::cout << "[main] arm comm thread started.\n";
    }
    else
    {
        std::cerr << "[main] sendRecvThread=null.\n";
    }

    arm.backToStart();
    arm.labelRun("forward");

    // TCP
    std::thread tcp_thread(tcpClient, std::ref(arm));
    std::cout << "[main] tcp client thread.\n";

    // keyboard
    std::thread keyboard_thread(monitorKeyboard, std::ref(exit_flag));
    std::cout << "[main] kb thread.\n";

    tcp_thread.join();
    std::cout << "[main] tcp thread joined.\n";

    auto t = std::time(nullptr);
    auto tmptr = std::localtime(&t);
    char buf[128];
    if (std::strftime(buf, sizeof(buf), "arm_position_log_%Y%m%d_%H%M%S.txt", tmptr) == 0)
    {
        std::cerr << "[main] strftime fail.\n";
        exit_flag = true;
    }
    std::string filename(buf);
    std::cout << "[main] log file= " << filename << "\n";

    std::thread log_thread;
    if (!exit_flag)
    {
        log_thread = std::thread(logArmPosition, std::ref(arm), std::ref(exit_flag), filename);
        std::cout << "[main] logging thread.\n";
    }

    std::thread arm_ctrl_thread;
    std::thread receive_thread;
    std::thread control_thread;
    std::thread serial_thread;
    std::thread force_send_thread;

    if (arm.tcp_sockfd != -1)
    {
        arm_ctrl_thread = std::thread(&Z1ARM::armCtrlByFSM, &arm);
        std::cout << "[main] armCtrlByFSM thread.\n";

        // 启动只处理最新包的接收线程
        receive_thread = std::thread(receiveDataFromServer, std::ref(arm));
        std::cout << "[main] recv thread.\n";

        if (serial_fd >= 0)
        {
            serial_thread = std::thread(readSerialData, serial_fd, std::ref(exit_flag));
            std::cout << "[main] serial thread.\n";
        }

        force_send_thread = std::thread(sendForceThreadFunc, std::ref(arm));
        std::cout << "[main] force send.\n";

#ifdef USE_POSITION_CONTROL
        control_thread = std::thread(armControlThreadFunc, std::ref(arm));
        std::cout << "[main] position ctrl.\n";
#elif defined(USE_SPEED_CONTROL)
        control_thread = std::thread(armSpeedControlThreadFunc, std::ref(arm));
        std::cout << "[main] speed ctrl.\n";
#endif

        while (!exit_flag)
        {
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
        }
        std::cout << "[main] exit_flag set, shutdown...\n";

        close(arm.tcp_sockfd);
        arm.tcp_sockfd = -1;
        std::cout << "[main] tcp closed.\n";

        if (receive_thread.joinable())
        {
            receive_thread.join();
            std::cout << "[main] recv joined.\n";
        }

        if (control_thread.joinable())
        {
            control_thread.join();
            std::cout << "[main] ctrl joined.\n";
        }

        if (arm_ctrl_thread.joinable())
        {
            arm_ctrl_thread.join();
            std::cout << "[main] armCtrlByFSM joined.\n";
        }
    }
    else
    {
        std::cerr << "[main] cannot connect server.\n";
        while (!exit_flag)
        {
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
        }
    }

    force_cv.notify_all();
    if (force_send_thread.joinable())
    {
        force_send_thread.join();
        std::cout << "[main] force send joined.\n";
    }

    if (keyboard_thread.joinable())
    {
        keyboard_thread.join();
        std::cout << "[main] kb joined.\n";
    }

    if (log_thread.joinable())
    {
        log_thread.join();
        std::cout << "[main] logging joined.\n";
    }
    if (serial_thread.joinable())
    {
        serial_thread.join();
        std::cout << "[main] serial joined.\n";
    }

    std::cout << "[main] restore arm.\n";
    arm.backToStart();
    arm.setFsm(ArmFSMState::PASSIVE);
    std::cout << "[main] arm restored.\n";

    if (arm.sendRecvThread)
    {
        arm.sendRecvThread->shutdown();
        std::cout << "[main] arm comm thread shut.\n";
    }

    std::cout << "[main] done.\n";
    return 0;
}